//
// Created by hamlet on 25-2-7.
//

#ifndef GPT2_H
#define GPT2_H

#include <random>

// 计算大于等于n的最小的8的倍数
constexpr int up8(const int n) {
    return n + 7 & ~7;
}

// ----------------------------------------------------------------------------
// GPT-2 model definition

/**
 * @brief GPT-2模型配置结构体
 */
struct GPT2Config {
    int max_seq_len;       // 最大序列长度，例如1024
    int vocab_size;        // 词表大小，例如50257
    int padded_vocab_size; // 填充后的词表大小，填充至128的倍数，例如50304
    int num_layers;        // 层数，例如12
    int num_heads;         // 注意力头数，例如12
    int channels;          // 隐藏层维度，例如768
};

/**
 * @brief 模型参数张量结构体
 * 存储模型所有可训练参数
 */
struct ParameterTensors {
    static constexpr int NUM_TENSORS = 16;

    float* wte;      // token嵌入权重 (V, C)
    float* wpe;      // 位置嵌入权重 (maxT, C)
    float* ln1w;     // 第一层归一化权重 (L, C)
    float* ln1b;     // 第一层归一化偏置 (L, C)
    float* qkvw;     // QKV投影权重 (L, 3*C, C)
    float* qkvb;     // QKV投影偏置 (L, 3*C)
    float* attprojw; // 注意力输出投影权重 (L, C, C)
    float* attprojb; // 注意力输出投影偏置 (L, C)
    float* ln2w;     // 第二层归一化权重 (L, C)
    float* ln2b;     // 第二层归一化偏置 (L, C)
    float* fcw;      // 前馈网络权重 (L, 4*C, C)
    float* fcb;      // 前馈网络偏置 (L, 4*C)
    float* fcprojw;  // 前馈网络输出投影权重 (L, C, 4*C)
    float* fcprojb;  // 前馈网络输出投影偏置 (L, C)
    float* lnfw;     // 最终层归一化权重 (C)
    float* lnfb;     // 最终层归一化偏置 (C)

    float* memory; // 参数内存
    size_t size;   // 参数内存大小

    /**
     * @brief 构造函数：分配内存并设置指针
     * @param config 模型配置
     * @return 分配的内存起始地址
     */
    explicit ParameterTensors(const GPT2Config& config);
    ~ParameterTensors(); // 析构函数

private:
    /**
     * @brief 计算每个参数张量的大小
     * @param sizes 存储参数大小的数组
     * @param config 模型配置
     */
    static void fill_in_sizes(size_t* sizes, const GPT2Config& config);
};

/**
 * @brief 模型激活值张量结构体
 * 存储前向传播过程中的中间结果
 */
struct ActivationTensors {
    static constexpr int NUM_TENSORS = 23;

    float* encoded;   // 输入编码 (B, T, C)
    float* ln1;       // 第一层归一化输出 (L, B, T, C)
    float* ln1_mean;  // 第一层归一化均值 (L, B, T)
    float* ln1_rstd;  // 第一层归一化标准差倒数 (L, B, T)
    float* qkv;       // QKV投影输出 (L, B, T, 3*C)
    float* atty;      // 注意力输出 (L, B, T, C)
    float* preatt;    // 注意力分数 (L, B, NH, T, T)
    float* att;       // 注意力权重 (L, B, NH, T, T)
    float* attproj;   // 注意力输出投影 (L, B, T, C)
    float* residual2; // 第二个残差连接 (L, B, T, C)
    float* ln2;       // 第二层归一化输出 (L, B, T, C)
    float* ln2_mean;  // 第二层归一化均值 (L, B, T)
    float* ln2_rstd;  // 第二层归一化标准差倒数 (L, B, T)
    float* fch;       // 前馈网络隐藏层 (L, B, T, 4*C)
    float* fch_gelu;  // GELU激活后的隐藏层 (L, B, T, 4*C)
    float* fcproj;    // 前馈网络输出 (L, B, T, C)
    float* residual3; // 第三个残差连接 (L, B, T, C)
    float* lnf;       // 最终层归一化输出 (B, T, C)
    float* lnf_mean;  // 最终层归一化均值 (B, T)
    float* lnf_rstd;  // 最终层归一化标准差倒数 (B, T)
    float* logits;    // 输出logits (B, T, V)
    float* probs;     // 输出概率 (B, T, V)
    float* losses;    // 损失值 (B, T)

    float* memory; // 激活内存

    /**
     * @brief 构造函数：分配内存并设置指针
     * @param config 模型配置
     * @param B 批次大小
     * @param T 序列长度
     * @return 分配的内存起始地址
     */
    explicit ActivationTensors(const GPT2Config& config, int B, int T);
    ~ActivationTensors();

private:
    /**
     * @brief 计算每个激活张量的大小
     * @param sizes 存储激活大小的数组
     * @param config 模型配置
     * @param B 批次大小
     * @param T 序列长度
     */
    static void fill_in_sizes(size_t* sizes, const GPT2Config& config, int B, int T);
};

/**
 * @brief 采样器类
 * 用于从概率分布中采样下一个token
 */
class Sampler {
public:
    Sampler();
    /**
     * @brief 从概率分布中采样
     * @param probs 概率分布
     * @param size 分布大小
     * @return 采样得到的token索引
     */
    int operator()(const float* probs, int size);

private:
    std::random_device rd;
    std::mt19937 gen;
    std::discrete_distribution<> dist;
};

/**
 * @brief GPT-2模型类
 * 实现模型的前向传播和文本生成
 */
class GPT2 {
    GPT2Config config;        // 模型配置
    ParameterTensors* params; // 模型参数
    ActivationTensors* acts;  // 模型激活值
    int batch_size;           // 当前前向传播的批次大小
    int seq_len;              // 当前前向传播的序列长度
    int* inputs;              // 当前前向传播的输入tokens

public:
    /**
     * @brief 构造函数
     * @param checkpoint_path 模型检查点路径
     */
    explicit GPT2(const char* checkpoint_path);
    ~GPT2(); // 析构函数

    /**
     * @brief 前向传播
     * @param inputs 输入tokens
     * @param B 批次大小
     * @param T 序列长度
     */
    void forward(const int* inputs, int B, int T);

    /**
     * @brief 生成一个token
     * @param inputs 输入tokens
     * @param t 当前位置
     * @return 生成的token
     */
    int generate_one(const int* inputs, int t);

private:
    Sampler sampler;
};

#endif //GPT2_H
